package me.ygy.agent;
import com.sun.tools.attach.*;
import jdk.internal.org.objectweb.asm.*;
import java.lang.instrument.*;
import java.util.HashMap;
import static jdk.internal.org.objectweb.asm.Opcodes.*;
public class App {
    public static void main(String[] args) throws Exception {
        VirtualMachine vm = VirtualMachine.attach(args[0]); // pid
        vm.loadAgent("/Users/yuguangyuan/code/gitee/maxwell/src/main/java/classes/agent.jar"); //agent path
    }
    public static void premain(String agentArgs, Instrumentation inst) {}
    public static void agentmain(String agentArgs, Instrumentation inst) throws Exception {
        String className = "me.ygy.agent.TestClass"; // 要trace的类名
        String methodName = "testMethod"; // 要trace的方法名
        inst.addTransformer(new TraceTransformer(className, methodName), true);
        inst.retransformClasses(Class.forName(className));
    }
    public static class TraceTransformer implements ClassFileTransformer {
        public String className_i, methodName_i;
        public TraceTransformer(String cn , String mn){
            className_i = cn; methodName_i = mn;
        }
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, java.security.ProtectionDomain protectionDomain, byte[] classfileBuffer) {
            if(className_i.equals(className.replace("/", "."))){
                ClassReader reader = new ClassReader(classfileBuffer);
                ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_MAXS);
                reader.accept(new TraceClassVisitor(writer,methodName_i), ClassReader.SKIP_FRAMES);
                return writer.toByteArray();
            }
            return classfileBuffer;
        }
    }

    public static class TraceClassVisitor extends ClassVisitor {
        public String methodName ;
        public TraceClassVisitor(final ClassVisitor cv, String mn) {
            super(ASM5, cv);
            methodName = mn;
        }
        public MethodVisitor visitMethod(final int access, final String name, final String desc, final String signature, final String[] exceptions) {
            MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
            return name.equals(methodName) ? new TraceMethodVisitor(mv,name) : mv;
        }
    }
    public static class TraceMethodVisitor extends MethodVisitor {
        public String name ;
        public int preCode = -1, lastCode = -1 ;
        protected TraceMethodVisitor(MethodVisitor mv, String n) {
            super(Opcodes.ASM5, mv);
            name = n;
        }
        public void visitCode() { //方法体内开始时调用
            println(name);
            setStartTime(name);
            super.visitCode();
        }
        public void visitInsn(int opcode) {
            if (opcode <= Opcodes.RETURN && opcode >= Opcodes.IRETURN) { //输出最后一个方法的运行时间
                setEndTime(String.valueOf(lastCode));
                computeAndPrintTime(String.valueOf(lastCode));
                println("almost time cost：");
                setEndTime(name);
                computeAndPrintTime(name); //输出整体方法的时间
            }
            super.visitInsn(opcode);
        }
        public void visitMethodInsn(final int opcode, final String owner, final String name, final String descriptor, final boolean isInterface) { // SetEndTime , 第一次过滤掉，因为是没有方法体的
            if(preCode!=-1) {
                setEndTime(String.valueOf(preCode));
                computeAndPrintTime(String.valueOf(preCode)); // print Cost time
            }
            // 输出当前方法名
            println("-|"+owner.replace("/","")+"@"+name);
            setStartTime(String.valueOf(opcode)); // 当前方法set进去
            preCode = lastCode = opcode;
            super.visitMethodInsn(opcode,owner,name,descriptor,isInterface);
        }
        private void println(String obj) {
            mv.visitFieldInsn(GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
            mv.visitLdcInsn(obj);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V", false);
        }
        private void setStartTime(String key) {
            mv.visitLdcInsn(key);
            mv.visitMethodInsn(INVOKESTATIC, "java/lang/System", "nanoTime", "()J", false);
            mv.visitMethodInsn(INVOKESTATIC, "me/ygy/agent/App", "setStartTime", "(Ljava/lang/String;J)V", false);
        }
        private void setEndTime(String key) {
            mv.visitLdcInsn(key);
            mv.visitMethodInsn(INVOKESTATIC, "java/lang/System", "nanoTime", "()J", false);
            mv.visitMethodInsn(INVOKESTATIC, "me/ygy/agent/App", "setEndTime", "(Ljava/lang/String;J)V", false);
        }

        private void computeAndPrintTime(String key) {
            mv.visitFieldInsn(GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
            mv.visitLdcInsn(key);
            mv.visitMethodInsn(INVOKESTATIC, "me/ygy/agent/App", "getCostTime", "(Ljava/lang/String;)Ljava/lang/String;", false);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V", false);
        }
    }
    public static HashMap<String, Long> sStartTime = new HashMap<>(), sEndTime = new HashMap<>();
    public static void setStartTime(String methodName, long time) {
        sStartTime.put(methodName, time);
    }
    public static void setEndTime(String methodName, long time) {
        sEndTime.put(methodName, time);
    }
    public static String getCostTime(String methodName) {
        return "-1".equals(methodName) ? "" : "--|method  Cost: " + (Long.valueOf(sEndTime.get(methodName) - sStartTime.get(methodName))*0.000001)+ " ms";
    }
}